package com.ptp.backendcertms.basiccompnents.thread;

import com.ptp.backendcertms.basiccompnents.log.TraceUtil;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;

import java.util.concurrent.*;

/**
 * 自定义线程池
 *
 * @author songcongcong
 * @date 2025/5/1
 */
@Slf4j
public class MyThreadPoolExecutor extends ThreadPoolExecutor {

    /**
     * 构造自定义线程池
     */
    public MyThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory);
    }

    @Override
    public void execute(@NotNull Runnable task) {
        super.execute(this.wrapRunnable(task));
    }

    @NotNull
    @Override
    public <T> Future<T> submit(@NotNull Runnable task, T result) {
        return super.submit(this.wrapRunnable(task), result);
    }

    @NotNull
    @Override
    public <T> Future<T> submit(@NotNull Callable<T> task) {
        return super.submit(this.wrapCallable(task));
    }

    @NotNull
    @Override
    public Future<?> submit(@NotNull Runnable task) {
        return super.submit(this.wrapRunnable(task));
    }

    /**
     * 包装传入tranceId
     */
    private Runnable wrapRunnable(final Runnable runnable) {
        String traceId = TraceUtil.getTraceId();
        return () -> {
            TraceUtil.setSubTraceId(traceId);
            try {
                runnable.run();
            } catch (Exception e) {
                log.error("线程任务runnable执行异常", e);
                throw e;
            }
        };
    }

    /**
     * 包装传入tranceId
     */
    private <T> Callable<T> wrapCallable(final Callable<T> callable) {
        String traceId = TraceUtil.getTraceId();
        return () -> {
            TraceUtil.setSubTraceId(traceId);
            try {
                return callable.call();
            } catch (Exception e) {
                log.error("线程任务callable执行异常", e);
                throw e;
            }
        };

    }
}
